import torch
import torch.nn as nn
import numpy as np
import time, sys, os
import copy
import random
import colorsys
import trimesh
from collections import namedtuple
from lib.fk.utils import get_4x4_matrix

ShadowHandOutput = namedtuple(
    "ShadowHandOutput",
    [
        "verts",
        "key_pts_pos",
    ],
)
ShadowHandOutput.__new__.__defaults__ = (None,) * len(ShadowHandOutput._fields)

def q2R(q):
    """
    q: 四元数, q = [x, y, z, w]
    R: 旋转矩阵,左乘
    """
    q = q[[3,0,1,2]]  # (x,y,z,w) -> (w,x,y,z)
    R = np.zeros((3, 3), dtype=float)

    R[0, 0] = 1 - 2 * (q[2] ** 2 + q[3] ** 2)
    R[1, 1] = 1 - 2 * (q[1] ** 2 + q[3] ** 2)
    R[2, 2] = 1 - 2 * (q[1] ** 2 + q[2] ** 2)

    R[0, 1] = 2 * (q[1] * q[2] - q[0] * q[3])
    R[1, 0] = 2 * (q[1] * q[2] + q[0] * q[3])

    R[0, 2] = 2 * (q[1] * q[3] + q[0] * q[2])
    R[2, 0] = 2 * (q[1] * q[3] - q[0] * q[2])

    R[1, 2] = 2 * (q[2] * q[3] - q[0] * q[1])
    R[2, 1] = 2 * (q[2] * q[3] + q[0] * q[1])

    return R

class Rodrigues:
    def __init__(self, rot_axis_list, device):
        SHADOW_HAND_DOF_NUM = 22
        self.device = device
        rot_axis = torch.FloatTensor(rot_axis_list).to(self.device).view(22, 3, 1)   # [22, 3, 1]
        rot_axis_T = torch.transpose(rot_axis, 1, 2) # [22, 1, 3]
        self.nnT = torch.matmul(rot_axis, rot_axis_T)  # [22, 3, 3]
        self.skew = torch.zeros_like(self.nnT).to(self.device) # [22, 3, 3]
        self.skew[:, 0, 1] = -rot_axis[:, 2, 0]
        self.skew[:, 0, 2] = rot_axis[:, 1, 0]
        self.skew[:, 1, 0] = rot_axis[:, 2, 0]
        self.skew[:, 1, 2] = -rot_axis[:, 0, 0]
        self.skew[:, 2, 0] = -rot_axis[:, 1, 0]
        self.skew[:, 2, 1] = rot_axis[:, 0, 0]

    '''
    description: 
    param {*} self
    param {*} theta : [batch, 22] 关节角度(rad)
    return {*} mat: [batch, 22, 3, 3]
    '''
    def rot_axis_to_mat(self, theta):
        batch_size = theta.size()[0]
        cos_theta = torch.cos(theta).view(batch_size, 22, 1, 1).repeat(1, 1, 3, 3)  # [batch, 22, 3, 3]
        sin_theta = torch.sin(theta).view(batch_size, 22, 1, 1).repeat(1, 1, 3, 3)  # [batch, 22, 3, 3]
        I = torch.eye(3).to(self.device).view(1, 1, 3, 3).repeat(batch_size, 22, 1, 1) # [batch, 22, 3, 3]
        nnT = self.nnT.unsqueeze(0).repeat(batch_size, 1, 1, 1) # [batch, 22, 3, 3]
        skew = self.skew.unsqueeze(0).repeat(batch_size, 1, 1, 1) # [batch, 22, 3, 3]
        mat = cos_theta * I + (1 - cos_theta) * nnT + sin_theta * skew
        return mat

class ShadowLayer(nn.Module):
    def __init__(self, device, sampled_pts_path, mesh_path):
        super(ShadowLayer, self).__init__()
        self.J = 22
        self.device = device

        self.body_name_list = ['palm',  # 1
                                   'knuckle', 'F3', 'F2', 'F1',  # 2, 3, 4, 5,
                                   'knuckle', 'F3', 'F2', 'F1',  # 6, 7, 8, 9,
                                   'knuckle', 'F3', 'F2', 'F1',  # 10, 11, 12, 13,
                                   'lfmetacarpal', 'knuckle', 'F3', 'F2', 'F1', # 14, 15, 16, 17, 18,
                                   'TH3_z', '', 'TH2_z', '', 'TH1_z',] # 19, 20, 21, 22, 23

        self.BODY_NUM = len(self.body_name_list) # 1 palm + 22 joints = 23

        TH3_z_q = [ 0, 0.3824995, 0, 0.9239557]
        body_t_list = [[0, 0, 0], 
                       [0.033, 0, 0.095], [0, 0, 0], [0, 0, 0.045], [0, 0, 0.025],
                       [0.011, 0, 0.099], [0, 0, 0], [0, 0, 0.045], [0, 0, 0.025],
                       [-0.011, 0, 0.095],[0, 0, 0], [0, 0, 0.045], [0, 0, 0.025],
                       [-0.034, 0, 0.021], [-0.001, 0, 0.067], [0, 0, 0], [0, 0, 0.045], [0, 0, 0.025],
                       [0.034, -0.009, 0.029], [0, 0, 0], [0, 0, 0.038], [0, 0, 0], [0, 0, 0.032]]
        M_sh = []  # Local transform matrix for shadow hand
        for i in range(1, self.BODY_NUM):
            if self.body_name_list[i] == 'TH3_z':
                local = get_4x4_matrix(TH3_z_q, body_t_list[i]).astype(np.float32)
            else:
                local = get_4x4_matrix([0,0,0,1], body_t_list[i]).astype(np.float32)
            M_sh.append(local)
        self.M_sh = torch.from_numpy(np.array(M_sh, dtype=np.float32)).to(self.device) # [22, 4, 4]
                
        rot_axis_list = [[0, 1, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0],
                        [0, 1, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0],
                        [0, 1, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0],
                        [0.571, 0, 0.821], [0, 1, 0], [1, 0, 0], [1, 0, 0], [1, 0, 0],
                        [0, 0, -1], [1, 0, 0], [1, 0, 0], [0, 1, 0], [0, 1, 0]]
        self.rodrigues = Rodrigues(rot_axis_list, device=self.device)

        self.parents = [-1,
                    0, 1, 2, 3,
                    0, 5, 6, 7,
                    0, 9, 10, 11,
                    0, 13, 14, 15, 16,
                    0, 18, 19, 20, 21]

        self.pts_dict = np.load(sampled_pts_path, allow_pickle=True).item()
        for name in self.pts_dict:
            self.pts_dict[name] = torch.from_numpy(self.pts_dict[name]).float().to(self.device)
        
        self.mesh_root = mesh_path
        
        # hand_info = np.load('/media/wind/Share/Projects/graspd/grasping/data/assets_shadow/unity_shadow.npz')
        # self.vert_num = hand_info['verts'].shape[0]        
        
        # self.faces = hand_info['faces']
        # self.vert_part_idxs = hand_info['vert_part_idxs']
        # verts = torch.from_numpy(hand_info['verts']).view(self.vert_num, 1, 3)
        
        # glob = self.transforms_blank(1).squeeze().detach().cpu()  # [J+1,4,4]
        # # begin from 1, [1, 23)
        # for i in range(1, len(self.parents)):
        #     # [bs, 4, 4]
        #     glob[i] = torch.matmul(glob[self.parents[i]], self.M_sh[i-1].detach().cpu())
        # T_verts_to_local = torch.eye(4).view(1, 4, 4).contiguous().repeat(self.vert_num, 1, 1)
        # T_verts_to_local = glob[self.vert_part_idxs]
        # verts = torch.matmul(verts - T_verts_to_local[:, :3, 3].unsqueeze(1), T_verts_to_local[:, :3, :3]).view(self.vert_num, 3)
        
        # import trimesh
        # for i in range(23):
        #     print(i)
        #     idx_set = np.where(self.vert_part_idxs == i)[0]
        #     if idx_set.shape[0] == 0:
        #         continue
        #     pc = trimesh.PointCloud(verts[idx_set])
        #     print(pc.vertices.shape)
        #     trimesh.Scene([pc]).show()
                    
        # self.verts = verts.to(self.device).view(1, self.vert_num, 1, 3).contiguous()
        # self.verts = verts.to(self.device).view(1, self.vert_num, 3).contiguous()
        
        # self.idxs = []
        # for i in range(23):
        #     idxs = np.where(self.vert_part_idxs == i)[0]
        #     self.idxs.append(idxs)
            
        self.theta_up_limit  = torch.FloatTensor([  0.349, 1.571, 1.571, 1.571,
                             0.349, 1.571, 1.571, 1.571, 
                             0.349, 1.571, 1.571, 1.571,
                             0.785, 0.349, 1.571, 1.571, 1.571,
                             1.047, 1.222, 0.209, 0.524, 0.000]).to(self.device)
        self.theta_low_limit = torch.FloatTensor([ -0.349, 0.000, 0.000, 0.000,
                            -0.349, 0.000, 0.000, 0.000,
                            -0.349, 0.000, 0.000, 0.000,
                             0.000,-0.349, 0.000, 0.000, 0.000,
                            -1.047, 0.000,-0.209,-0.524,-1.571]).to(self.device)
        self.theta_up_limit.requires_grad = False
        self.theta_low_limit.requires_grad = False
    
    def transforms_blank(self, bs, joints=23):
        """
        ts : (bs, 23, 4, 4) ndarray
            Array of identity transforms
        """
        diagonal = torch.eye(4).to(self.device)
        ts = diagonal.view(1, 1, 4, 4).repeat(bs, joints, 1, 1)
        # diagonal = torch.eye(4).to(self.device)
        # ts = diagonal.expand(bs, joints, 4, 4)        
        return ts

    def transforms_base(self, base):
        """
        base:(bs, 7)—— x y z qx qy qz qw
        """
        base = base.unsqueeze(-2)
        rotations = base[:, :, 3:]
        q_length = torch.sqrt(torch.sum(rotations.pow(2), dim=-1))  # [bs,J,1]
        qx = rotations[..., 0] / q_length  # [bs,J,1]
        qy = rotations[..., 1] / q_length  # [bs,J,1]
        qz = rotations[..., 2] / q_length  # [bs,J,1]
        qw = rotations[..., 3] / q_length  # [bs,J,1]
        """Unit quaternion based rotation matrix computation"""
        x2 = qx + qx  # [bs,J,1]
        y2 = qy + qy
        z2 = qz + qz
        xx = qx * x2
        yy = qy * y2
        wx = qw * x2
        xy = qx * y2
        yz = qy * z2
        wy = qw * y2
        xz = qx * z2
        zz = qz * z2
        wz = qw * z2

        dim0 = torch.stack([1.0 - (yy + zz), xy - wz, xz + wy], -1)
        dim1 = torch.stack([xy + wz, 1.0 - (xx + zz), yz - wx], -1)
        dim2 = torch.stack([xz - wy, yz + wx, 1.0 - (xx + yy)], -1)
        R = torch.stack([dim0, dim1, dim2], -2).to(self.device)   # [bs,1,3,3]

        T = base[..., :3].unsqueeze(-1).to(self.device)  # (bs, 1, 3, 1)
        zeros = torch.zeros([int(base.shape[0]), 1, 1, 3]).to(self.device)  # (bs, 1, 1, 3)
        ones = torch.ones([int(base.shape[0]), 1, 1, 1]).to(self.device)  # (bs, 1, 1, 1)
        base_M = torch.cat([torch.cat([R, zeros], -2), torch.cat([T, ones], -2)], -1) # (bs, 1, 4, 4)

        return base_M   # [F,1,4,4]

    def set_rotations_into_window(self, rotations):
        bs = rotations.size(0)
        max_ = self.theta_up_limit.to(self.device).view(1, -1).repeat(bs, 1)
        min_ = self.theta_low_limit.to(self.device).view(1, -1).repeat(bs, 1)
        rotations = torch.where(rotations > max_, max_, rotations)
        rotations = torch.where(rotations < min_, min_, rotations)
        return rotations
    
    def transforms_rotations(self, rotations):
        bs = rotations.size(0)
        rot = rotations.clone()
        full_val = torch.ones_like(rot[:, 0]) * 1.571 # 90 degree
        null_val = torch.zeros_like(rot[:, 0])
        tmp1 = rot[:, 2] + rot[:, 3]
        rot[:, 2] = torch.where(tmp1 > 1.571, full_val, tmp1)
        rot[:, 3] = torch.where(tmp1 > 1.571, tmp1 - 1.571, null_val)
        tmp2 = rot[:, 6] + rot[:, 7]
        rot[:, 6] = torch.where(tmp2 > 1.571, full_val, tmp2)
        rot[:, 7] = torch.where(tmp2 > 1.571, tmp2 - 1.571, null_val)
        tmp3 = rot[:, 10] + rot[:, 11]
        rot[:, 10] = torch.where(tmp3 > 1.571, full_val, tmp3)
        rot[:, 11] = torch.where(tmp3 > 1.571, tmp3 - 1.571, null_val)  
        tmp4 = rot[:, 15] + rot[:, 16]
        rot[:, 15] = torch.where(tmp4 > 1.571, full_val, tmp4)
        rot[:, 16] = torch.where(tmp4 > 1.571, tmp4 - 1.571, null_val)                         
        M_r = torch.eye(4).to(self.device).view(1, 1, 4, 4).contiguous().repeat(bs, self.J, 1, 1)
        M_r[:, :, :3, :3] = self.rodrigues.rot_axis_to_mat(rot)
        return M_r   # [bs,J,4,4]

    def joint2index(self, joints):
        """
        Food:        (0,  1),  2,     3  --            (FF4, FF3),   FF2,     FF1
        Middle:      (4,  5),  6,     7  --            (MF4, MF3),   MF2,     MF1
        Ring:        (8,  9), 10,    11  --            (RF4, RF3),   RF2,     RF1
        little: 12,  (13,14), 15,    16  --    LF5,    (LF4, LF3),   LF2,     LF1
        Thumb:       (17,18),(19,20),21  --            (TF5, TF4),(TF3, TF2), TF1
        """
                
        joints_1_index = [0, 1, 4, 5, 8, 9, 12, 13, 14, 17, 18]
        joints_2_index = [2, 6, 10, 15, 19, 20]
        joints_3_index = [3, 7, 11, 16, 21]
        
        index = torch.zeros([3, self.J], dtype=torch.bool).to(self.device)
        index[0, joints_1_index] = True
        index[1, joints_2_index] = True
        index[2, joints_3_index] = True
        if isinstance(joints, int):
            assert joints in [1, 2, 3], 'joints should be in [1, 2, 3]'
            return index[joints - 1]
        elif isinstance(joints, list):
            index_merged = torch.zeros(self.J, dtype=torch.bool).to(self.device)
            for j in joints:
                assert j in [1, 2, 3], 'joints should be in [1, 2, 3]'
                index_merged = index_merged | index[j - 1]
            return index_merged

    def forward(self, rotations):
        bs = rotations.size(0)
        base = torch.eye(4).view(1, 4, 4).repeat(bs, 1, 1).to(self.device)
        base.requires_grad = False
        # transforms_local
        M_r = self.transforms_rotations(rotations).to(self.device) # [bs,J,4,4]
        M_sh = self.M_sh.clone().unsqueeze(0).repeat(bs, 1, 1, 1)  # [bs,J,4,4]
        locals = torch.matmul(M_sh, M_r)  # [bs,J,4,4]
        
        glob_ff_0 = torch.matmul(base, locals[:,0])
        glob_mf_0 = torch.matmul(base, locals[:,4])
        glob_rf_0 = torch.matmul(base, locals[:,8])
        glob_lf_0 = torch.matmul(base, locals[:,12])
        glob_th_0 = torch.matmul(base, locals[:,17])
        
        glob_ff_1 = torch.matmul(glob_ff_0, locals[:,1])
        glob_mf_1 = torch.matmul(glob_mf_0, locals[:,5])
        glob_rf_1 = torch.matmul(glob_rf_0, locals[:,9])
        glob_lf_1 = torch.matmul(glob_lf_0, locals[:,13])
        glob_th_1 = torch.matmul(glob_th_0, locals[:,18])
        
        glob_ff_2 = torch.matmul(glob_ff_1, locals[:,2])
        glob_mf_2 = torch.matmul(glob_mf_1, locals[:,6])
        glob_rf_2 = torch.matmul(glob_rf_1, locals[:,10])
        glob_lf_2 = torch.matmul(glob_lf_1, locals[:,14])
        glob_th_2 = torch.matmul(glob_th_1, locals[:,19])

        glob_ff_3 = torch.matmul(glob_ff_2, locals[:,3])
        glob_mf_3 = torch.matmul(glob_mf_2, locals[:,7])
        glob_rf_3 = torch.matmul(glob_rf_2, locals[:,11])
        glob_lf_3 = torch.matmul(glob_lf_2, locals[:,15])
        glob_th_3 = torch.matmul(glob_th_2, locals[:,20])        
        
        glob_lf_4 = torch.matmul(glob_lf_3, locals[:,16])
        glob_th_4 = torch.matmul(glob_th_3, locals[:,21])
        
        globals = torch.cat([base.unsqueeze(1), 
                             glob_ff_0.unsqueeze(1),glob_ff_1.unsqueeze(1),glob_ff_2.unsqueeze(1),glob_ff_3.unsqueeze(1),
                             glob_mf_0.unsqueeze(1),glob_mf_1.unsqueeze(1),glob_mf_2.unsqueeze(1),glob_mf_3.unsqueeze(1),
                             glob_rf_0.unsqueeze(1),glob_rf_1.unsqueeze(1),glob_rf_2.unsqueeze(1),glob_rf_3.unsqueeze(1),
                             glob_lf_0.unsqueeze(1),glob_lf_1.unsqueeze(1),glob_lf_2.unsqueeze(1),glob_lf_3.unsqueeze(1),glob_lf_4.unsqueeze(1),
                             glob_th_0.unsqueeze(1),glob_th_1.unsqueeze(1),glob_th_2.unsqueeze(1),glob_th_3.unsqueeze(1),glob_th_4.unsqueeze(1)], dim=1)      

        key_points_pos = globals[:, :, :3, 3]
        
        # verts = self.verts.clone().repeat(bs, 1, 1)
        # verts.requires_grad = False
        # verts_rot = torch.zeros_like(verts).to(self.device)
        # for i in range(23):
        #     if (i == 18 or i == 20):
        #         continue
        #     # idxs = np.where(self.vert_part_idxs == i)[0]
        #     idxs = self.idxs[i]
        #     verts_rot[:, idxs] = torch.matmul(verts[:, idxs], globals[:, i, :3, :3].transpose(1, 2).contiguous()) + globals[:, i, :3, 3].unsqueeze(-2)
        
        transformed_finger_pts_dict = dict()
        transformed_pts_list = []
        transformed_tips_list = []
        for i in range(self.BODY_NUM):
            if self.body_name_list[i] in ['palm', 'lfmetacarpal', 'F3', 'F2', 'F1', 'TH3_z', 'TH2_z', 'TH1_z']:
                pts = self.pts_dict[self.body_name_list[i]].unsqueeze(0).repeat(bs, 1, 1) # [F, 100 / 400, 3]
                if i == 18 or i == 20:
                    t = globals[:, i+1] # [F, 4, 4]
                else:
                    t = globals[:, i]
                R_mat = t[:,:3,:3].transpose(1, 2).contiguous() # [F, 3, 3]
                t_mat = t[:,:3,3].unsqueeze(1) # [F, 1, 3]///
                pts_trans = torch.matmul(pts, R_mat) + t_mat # [F, 100 / 400, 3], 400 only for palm
                if self.body_name_list[i] in ['F1', 'TH1_z']:
                    transformed_tips_list.append(pts_trans[:, -1, :].unsqueeze(1))
                # else:
                #     transformed_pts_list.append(pts)
                transformed_pts_list.append(pts_trans)

        verts_rot = torch.cat(transformed_pts_list, axis=1) # [F, 2000, 3] or [F, 1500, 3] (remove F1 & TH1 link)
        transformed_tips = torch.cat(transformed_tips_list, axis=1)  # [F, 5, 3]
        key_points_pos = torch.cat([key_points_pos, transformed_tips], axis=1)      
        
        output = ShadowHandOutput(verts=verts_rot, key_pts_pos=key_points_pos)
        
        return output
    
    def get_posed_pc(self, base, rotations):
        bs = rotations.size(0)
        base = self.transforms_base(base).squeeze(1)
        # transforms_local
        M_r = self.transforms_rotations(rotations).to(self.device) # [bs,J,4,4]
        M_sh = self.M_sh.clone().unsqueeze(0).repeat(bs, 1, 1, 1)  # [bs,J,4,4]
        locals = torch.matmul(M_sh, M_r)  # [bs,J,4,4]
        
        glob_ff_0 = torch.matmul(base, locals[:,0])
        glob_mf_0 = torch.matmul(base, locals[:,4])
        glob_rf_0 = torch.matmul(base, locals[:,8])
        glob_lf_0 = torch.matmul(base, locals[:,12])
        glob_th_0 = torch.matmul(base, locals[:,17])
        
        glob_ff_1 = torch.matmul(glob_ff_0, locals[:,1])
        glob_mf_1 = torch.matmul(glob_mf_0, locals[:,5])
        glob_rf_1 = torch.matmul(glob_rf_0, locals[:,9])
        glob_lf_1 = torch.matmul(glob_lf_0, locals[:,13])
        glob_th_1 = torch.matmul(glob_th_0, locals[:,18])
        
        glob_ff_2 = torch.matmul(glob_ff_1, locals[:,2])
        glob_mf_2 = torch.matmul(glob_mf_1, locals[:,6])
        glob_rf_2 = torch.matmul(glob_rf_1, locals[:,10])
        glob_lf_2 = torch.matmul(glob_lf_1, locals[:,14])
        glob_th_2 = torch.matmul(glob_th_1, locals[:,19])

        glob_ff_3 = torch.matmul(glob_ff_2, locals[:,3])
        glob_mf_3 = torch.matmul(glob_mf_2, locals[:,7])
        glob_rf_3 = torch.matmul(glob_rf_2, locals[:,11])
        glob_lf_3 = torch.matmul(glob_lf_2, locals[:,15])
        glob_th_3 = torch.matmul(glob_th_2, locals[:,20])        
        
        glob_lf_4 = torch.matmul(glob_lf_3, locals[:,16])
        glob_th_4 = torch.matmul(glob_th_3, locals[:,21])
            
        globals = torch.cat([base.unsqueeze(1), 
                             glob_ff_0.unsqueeze(1),glob_ff_1.unsqueeze(1),glob_ff_2.unsqueeze(1),glob_ff_3.unsqueeze(1),
                             glob_mf_0.unsqueeze(1),glob_mf_1.unsqueeze(1),glob_mf_2.unsqueeze(1),glob_mf_3.unsqueeze(1),
                             glob_rf_0.unsqueeze(1),glob_rf_1.unsqueeze(1),glob_rf_2.unsqueeze(1),glob_rf_3.unsqueeze(1),
                             glob_lf_0.unsqueeze(1),glob_lf_1.unsqueeze(1),glob_lf_2.unsqueeze(1),glob_lf_3.unsqueeze(1),glob_lf_4.unsqueeze(1),
                             glob_th_0.unsqueeze(1),glob_th_1.unsqueeze(1),glob_th_2.unsqueeze(1),glob_th_3.unsqueeze(1),glob_th_4.unsqueeze(1)], dim=1)      


        key_points_pos = globals[:, :, :3, 3]

        transformed_pts_list = []
        transformed_tips_list = []
        for i in range(self.BODY_NUM):
            if self.body_name_list[i] in ['palm', 'lfmetacarpal', 'F3', 'F2', 'F1', 'TH3_z', 'TH2_z', 'TH1_z']:
                pts = self.pts_dict[self.body_name_list[i]].unsqueeze(0).repeat(bs, 1, 1) # [F, 100 / 400, 3]
                if i == 18 or i == 20:
                    t = globals[:, i+1] # [F, 4, 4]
                else:
                    t = globals[:, i]
                R_mat = t[:,:3,:3].transpose(1, 2).contiguous() # [F, 3, 3]
                t_mat = t[:,:3,3].unsqueeze(1) # [F, 1, 3]///
                pts_trans = torch.matmul(pts, R_mat) + t_mat # [F, 100 / 400, 3], 400 only for palm
                if self.body_name_list[i] in ['F1', 'TH1_z']:
                    transformed_tips_list.append(pts_trans[:, -1, :].unsqueeze(1))
                # else:
                #     transformed_pts_list.append(pts)
                transformed_pts_list.append(pts_trans)

        verts_rot = torch.cat(transformed_pts_list, axis=1) # [F, 2000, 3] or [F, 1500, 3] (remove F1 & TH1 link)
        transformed_tips = torch.cat(transformed_tips_list, axis=1)  # [F, 5, 3]
        key_points_pos = torch.cat([key_points_pos, transformed_tips], axis=1)      
        
        output = ShadowHandOutput(verts=verts_rot, key_pts_pos=key_points_pos)
        
        return output    
    
    def compute_loss_range(self, rotations):
        target_rot = (self.theta_up_limit + self.theta_low_limit) * 0.5
        # 小拇指第一个关节
        target_rot[12] = 0.1
        l_range = ((rotations - target_rot[None,:]) ** 2).sum()
        return l_range
    
    def compute_loss_limit(self, rotations):
        l_limit_1 = -torch.minimum(rotations - self.theta_low_limit[None, :], torch.zeros_like(rotations)).sum()
        l_limit_2 = -torch.minimum(self.theta_up_limit[None, :] - rotations, torch.zeros_like(rotations)).sum()
        return l_limit_1 + l_limit_2
    
    def compute_loss_self_collision(self, key_pts):
        def subset_dis(pts, idx1_list, idx2_list):
                subset_dis = ((pts[:, idx1_list].unsqueeze(2).expand(-1, -1, len(idx2_list), -1)
                            - pts[:, idx2_list].unsqueeze(1).expand(-1, len(idx1_list), -1, -1)) ** 2).sum(-1).sqrt().reshape(pts.shape[0], -1)
                return subset_dis
        # lines = [[0, 1],  [2, 3],   [3, 4],   [4, 23],               # food
        #          [0, 5],  [6, 7],   [7, 8],   [8, 24],               # middle
        #          [0, 9],  [10, 11], [11, 12], [12, 25],              # ring
        #[0, 13], [13, 14], [15, 16], [16, 17], [17, 26],    # little
        #          [0, 18], [19, 20], [21, 22], [22, 27]]              # thumb
        d1 = subset_dis(key_pts, [3, 4,], [7, 8,])
        d2 = subset_dis(key_pts, [7, 8,], [11, 12,])
        d3 = subset_dis(key_pts, [11, 12,], [16, 17,])
        d4 = subset_dis(key_pts, [15, 16, 17,], [2, 3, 4, 6, 7, 8, 10, 11, 12, 19, 21, 22,])
        d5 = subset_dis(key_pts, [21, 22,], [2, 3, 4, 6, 7, 8, 10, 11, 12, 15, 16, 17,])
        d = torch.cat([d1, d2, d3, d4, d5], dim=1) - 0.03
        loss = torch.where(d < 0, -d, torch.zeros_like(d)).sum()
        # loss = torch.mean(loss, dim=1, keepdim=False)
        return loss
    
    def get_posed_mesh(self, base, rotation):
        
        rotation = torch.from_numpy(rotation).to(self.device)
        # if (rotation < self.theta_low_limit).any():
        #     print('too small rot!')
        # if (rotation > self.theta_up_limit).any():
        #     print('too large rot!')
            
        rotation = rotation.view(1, 22)
        M_r = self.transforms_rotations(rotation).view(22, 4, 4).detach()
        M_sh = self.M_sh.clone().detach()
        locals = torch.matmul(M_sh, M_r).cpu().numpy()
        
        T = get_4x4_matrix(pos=base[:3], quat=base[3:])
        body_globals = np.eye(4)
        body_globals = body_globals[np.newaxis, :]
        body_globals = np.repeat(body_globals, 23, axis=0)
        for i in range(1, 23):
            body_globals[i] = np.matmul(body_globals[self.parents[i]], locals[i-1])

        mesh_list = []
        for i in range(23):
            if self.body_name_list[i] != '':
                mesh = trimesh.load(os.path.join(self.mesh_root, f'{self.body_name_list[i]}.stl')) # mm -> m
                mesh.apply_scale(0.001)
                if i == 18 or i == 20:
                    mesh.apply_transform(body_globals[i+1])
                else:
                    mesh.apply_transform(body_globals[i])
                mesh.apply_transform(T)
                mesh.visual.face_colors = np.array([200, 200, 200, 255])
                mesh_list.append(mesh)
        merged_mesh = self.merge_meshes(mesh_list)
                
        return merged_mesh
    
    def get_rotated_hand(self, rotation):

        rotation = rotation.view(1, 22)
        M_r = self.transforms_rotations(rotation).view(22, 4, 4).detach()
        M_sh = self.M_sh.clone().detach()
        locals = torch.matmul(M_sh, M_r).cpu().numpy()
        
        body_globals = np.eye(4)
        body_globals = body_globals[np.newaxis, :]
        body_globals = np.repeat(body_globals, 23, axis=0)
        for i in range(1, 23):
            body_globals[i] = np.matmul(body_globals[self.parents[i]], locals[i-1])

        mesh_list = []
        for i in range(23):
            if self.body_name_list[i] != '':
                mesh = trimesh.load(os.path.join(self.mesh_root, f'{self.body_name_list[i]}.stl')) # mm -> m
                mesh.apply_scale(0.001)
                if i == 18 or i == 20:
                    mesh.apply_transform(body_globals[i+1])
                else:
                    mesh.apply_transform(body_globals[i])
                mesh_list.append(mesh)
        merged_mesh = self.merge_meshes(mesh_list)
                
        return merged_mesh
    
    def merge_meshes(self, mesh_list):
        verts = copy.deepcopy(mesh_list[0].vertices)
        faces = copy.deepcopy(mesh_list[0].faces)
        vert_num = verts.shape[0]
        for mesh in mesh_list[1:]:
            verts = np.vstack([verts, mesh.vertices])
            faces = np.vstack([faces, mesh.faces + vert_num])
            vert_num += mesh.vertices.shape[0]
            
        return trimesh.base.Trimesh(vertices=verts, faces=faces)
        
def get_n_hls_colors(num):
    hls_colors = []
    i = 0
    step = 360.0 / num
    while i < 360:
        h = i
        s = 90 + random.random() * 10
        l = 50 + random.random() * 10
        _hlsc = [h / 360.0, l / 100.0, s / 100.0]
        hls_colors.append(_hlsc)
        i += step
 
    return hls_colors
 
def ncolors(num):
    rgb_colors = []
    if num < 1:
        return rgb_colors
    hls_colors = get_n_hls_colors(num)
    for hlsc in hls_colors:
        _r, _g, _b = colorsys.hls_to_rgb(hlsc[0], hlsc[1], hlsc[2])
        r, g, b = [int(x * 255.0) for x in (_r, _g, _b)]
        rgb_colors.append([r, g, b])
 
    return rgb_colors
    
if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    rotations = [0.1, 0.5, 0.5, 0.5,
                    0, 0, 0, 0,
                    0, 0, 0, 0,
                 0.5, 0, 0, 0, 0,
                 0, 0, 0, 0, -1.0]
                
    rotations = torch.FloatTensor(rotations).view(1, -1).to(device)
    shlayer = ShadowLayer(device)

    print("Forward start")
    t_preforward = time.time()
    output = shlayer(rotations)
    print("Forward end, time: {:.2f}s".format(time.time() - t_preforward))
    print(output.verts.size())
    
    v = output.verts[0].detach().cpu().numpy()
    faces = shlayer.faces#.numpy()
    idxs = shlayer.vert_part_idxs#.numpy()
    colors = ncolors(23)
    import trimesh
    m = trimesh.base.Trimesh(vertices=v, faces=faces)
    for i in range(v.shape[0]):
        m.visual.vertex_colors[i][:3] = colors[idxs[i]]
        
    trimesh.Scene([m]).show()